//
//  AESUtils.m
//  XLCryptoTool
//
//  Created by ilongge on 2021/7/29.
//

#import "XLAESUtils.h"
#import <CommonCrypto/CommonCrypto.h>
#import <CommonCrypto/CommonDigest.h>

@implementation XLAESUtils

+ (NSString *)encryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andText:(NSString *)cryptText
{
    XLAESUtils *utils = [[XLAESUtils alloc] init];
    return [utils encryptWithKey:cryptKey andIvKey:ivKey andMode:mode andText:cryptText];
}
- (NSString *)encryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andText:(NSString *)cryptText
{
    if (cryptKey == nil) {
        return nil;
    }
    
    if (mode == XL_AES_Crypto_Mode_CBC && ivKey == nil) {
        return nil;
    }
    cryptKey = [XLAESUtils checkAESKey:cryptKey];
    if (mode == XL_AES_Crypto_Mode_CBC) {
        ivKey = [XLAESUtils checkAESKey:ivKey];
    }
    NSData *data = [cryptText dataUsingEncoding:NSUTF8StringEncoding];
    NSData *result = [XLAESUtils encryptWithKey:cryptKey andIvKey:ivKey andMode:mode andData:data];
    if (result && result.length > 0) {
        Byte *datas = (Byte*)[result bytes];
        NSMutableString *output = [NSMutableString stringWithCapacity:result.length * 2];
        for(int i = 0; i < result.length; i++){
            [output appendFormat:@"%02x", datas[i]];
        }
        return output;
    }
    return nil;
}

+ (NSData *)encryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andData:(NSData *)cryptData
{
    XLAESUtils *utils = [[XLAESUtils alloc] init];
    return [utils encryptWithKey:cryptKey andIvKey:ivKey andMode:mode andData:cryptData];
}

- (NSData *)encryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andData:(NSData *)cryptData
{
    cryptKey = [XLAESUtils checkAESKey:cryptKey];
    NSData *initVector;
    if (mode == XL_AES_Crypto_Mode_CBC) {
        ivKey = [XLAESUtils checkAESKey:ivKey];
        initVector = [ivKey dataUsingEncoding:NSUTF8StringEncoding];
    }
    else{
        initVector = nil;
    }
    char keyPtr[kCCKeySizeAES128+1];
    memset(keyPtr, 0, sizeof(keyPtr));
    [cryptKey getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
    NSUInteger dataLength = cryptData.length;
    size_t bufferSize = dataLength + kCCBlockSizeAES128;
    void *buffer = malloc(bufferSize);
    size_t numBytesEncrypted = 0;
    CCOptions options = [self optionsFor:mode];
    CCCryptorStatus cryptStatus = CCCrypt(kCCEncrypt,
                                          kCCAlgorithmAES,
                                          options,
                                          keyPtr,
                                          kCCBlockSizeAES128,
                                          initVector.bytes,
                                          cryptData.bytes,
                                          dataLength,
                                          buffer,
                                          bufferSize,
                                          &numBytesEncrypted);
    if (cryptStatus == kCCSuccess) {
        NSData *data = [NSData dataWithBytesNoCopy:buffer length:numBytesEncrypted];
        return data;
    }
    free(buffer);
    return nil;
}

+ (NSString *)decryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andText:(NSString *)cryptText
{
    XLAESUtils *utils = [[XLAESUtils alloc] init];
    return [utils decryptWithKey:cryptKey andIvKey:ivKey andMode:mode andText:cryptText];
}
- (NSString *)decryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andText:(NSString *)cryptText
{
    if (cryptKey == nil) {
        return nil;
    }
    
    if (mode == XL_AES_Crypto_Mode_CBC && ivKey == nil) {
        return nil;
    }
    
    cryptKey = [XLAESUtils checkAESKey:cryptKey];
    if (mode == XL_AES_Crypto_Mode_CBC) {
        ivKey = [XLAESUtils checkAESKey:ivKey];
    }
    NSMutableData *data = [NSMutableData dataWithCapacity:cryptText.length / 2];
    unsigned char whole_byte;
    char byte_chars[3] = {'\0','\0','\0'};
    for (int i = 0; i < [cryptText length] / 2; i++) {
        byte_chars[0] = [cryptText characterAtIndex:i * 2];
        byte_chars[1] = [cryptText characterAtIndex:i * 2 + 1];
        whole_byte = strtol(byte_chars, NULL, 16);
        [data appendBytes:&whole_byte length:1];
    }
    NSData *result = [XLAESUtils decryptWithKey:cryptKey andIvKey:ivKey andMode:mode andData:data];
    if (result && result.length > 0) {
        NSString *output = [[NSString alloc] initWithData:result encoding:NSUTF8StringEncoding];
        return output;
    }
    return nil;
}

+ (NSData *)decryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andData:(NSData *)cryptData
{
    XLAESUtils *utils = [[XLAESUtils alloc] init];
    return [utils decryptWithKey:cryptKey andIvKey:ivKey andMode:mode andData:cryptData];
}

- (NSData *)decryptWithKey:(NSString *)cryptKey andIvKey:(NSString *)ivKey andMode:(XL_AES_Crypto_Mode)mode andData:(NSData *)cryptData
{
    cryptKey = [XLAESUtils checkAESKey:cryptKey];
    NSData *initVector;
    if (mode == XL_AES_Crypto_Mode_CBC) {
        ivKey = [XLAESUtils checkAESKey:ivKey];
        initVector = [ivKey dataUsingEncoding:NSUTF8StringEncoding];
    }
    else{
        initVector = nil;
    }
    char keyPtr[kCCKeySizeAES128+1];
    memset(keyPtr, 0, sizeof(keyPtr));
    [cryptKey getCString:keyPtr maxLength:sizeof(keyPtr) encoding:NSUTF8StringEncoding];
    NSUInteger dataLength = cryptData.length;
    size_t bufferSize = dataLength + kCCBlockSizeAES128;
    void *buffer = malloc(bufferSize);
    CCOptions options = [self optionsFor:mode];
    size_t numBytesDecrypted = 0;
    CCCryptorStatus cryptStatus = CCCrypt(kCCDecrypt,
                                          kCCAlgorithmAES,
                                          options,
                                          keyPtr,
                                          kCCBlockSizeAES128,
                                          initVector.bytes,
                                          cryptData.bytes,
                                          dataLength,
                                          buffer,
                                          bufferSize,
                                          &numBytesDecrypted);
    if (cryptStatus == kCCSuccess) {
        NSData *data = [NSData dataWithBytesNoCopy:buffer length:numBytesDecrypted];
        return data;
    }
    free(buffer);
    return nil;
}


- (CCOptions)optionsFor:(XL_AES_Crypto_Mode)mode {
    CCOptions options;
    switch (mode) {
        case XL_AES_Crypto_Mode_CBC:
            options = kCCOptionPKCS7Padding;
            break;
        case XL_AES_Crypto_Mode_ECB:
            options = kCCOptionPKCS7Padding | kCCOptionECBMode;
            break;
        default:
            options = kCCOptionPKCS7Padding;
            break;
    }
    return options;
}

+ (NSString *)checkAESKey:(NSString *)AES_Key
{
    if (AES_Key == nil) {
        AES_Key = @"";
    }
    if (AES_Key.length < 16) {
        NSMutableArray *array = [NSMutableArray arrayWithArray:@[AES_Key]];
        while (array.count <= 16 - AES_Key.length) {
            [array insertObject:@"0" atIndex:0];
        }
        AES_Key = [array componentsJoinedByString:@""];
    }
    return AES_Key;
}
@end

